from qkdKey import QKDKey

class QKDBuffer:
    def __init__(self, key_size, max_capacity, min_capacity, ip_address):
        self.key_size = key_size  # Standard size for keys stored in buffer
        self.max_capacity = max_capacity  # Maximum storage depth Mmax
        self.min_capacity = min_capacity  # Minimum amount of key material required Mmin
        self.current_capacity = 0  # Current value Mcur(t)
        self.threshold = max_capacity // 2  # Threshold value Mthr(t), default to half of max capacity
        self.buffer = {}  # Storage for keys using a dictionary with key_id as identifier
        self.ip_address = ip_address  # Unique IP address to identify the buffer

    def add_key(self, key):
        if not isinstance(key, QKDKey):
            print("Invalid key type. Only QKDKey instances can be added to the buffer.")
            return
        if len(self.buffer) < self.max_capacity:
            if key.key_id not in self.buffer:
                self.buffer[key.key_id] = key
                self.current_capacity += 1
                print(f"Key added to buffer. Current capacity: {self.current_capacity}/{self.max_capacity}")
            else:
                print("Key with this ID already exists in the buffer.")
        else:
            print("Buffer is full. Cannot add more keys.")

    def get_key(self, key_id):
        if key_id in self.buffer:
            return self.buffer[key_id]
        else:
            print("Key not found in buffer.")
            return None

    def get_remove_key(self, key_id):
        if key_id in self.buffer:
            key = self.buffer.pop(key_id)
            self.current_capacity -= 1
            print(f"Key retrieved and removed from buffer. Current capacity: {self.current_capacity}/{self.max_capacity}")
            return key
        else:
            print("Key not found in buffer.")
            return None

    def generate_and_add_key(self):
        if len(self.buffer) < self.max_capacity:
            key = QKDKey.generate_random_key(self.key_size)
            self.add_key(key)
        else:
            print("Buffer is full. Cannot generate and add more keys.")

    def get_status(self):
        if self.current_capacity >= self.threshold:
            return "READY"
        elif self.threshold > self.current_capacity > self.min_capacity:
            return "WARNING"
        elif self.current_capacity <= self.min_capacity and len(self.buffer) > 0:
            return "CHARGING"
        elif self.current_capacity <= self.min_capacity and len(self.buffer) == 0:
            return "EMPTY"
        else:
            return "UNKNOWN"

    def monitor_buffer(self):
        status = self.get_status()
        print(f"QKD Buffer status for IP address {self.ip_address}: {status}")
        if status == "EMPTY":
            print("Priority action required: Buffer is EMPTY.")

    def print_all_keys(self):
        if not self.buffer:
            print("Buffer is empty. No keys to display.")
        else:
            print(f"Current keys in buffer for IP address {self.ip_address}:")
            for key_id, key in self.buffer.items():
                print(f"Key ID: {key_id}, Key Value: {key.key_value.hex()}")

if __name__ == "__main__":
    # Example usage of QKDBuffer
    buffer1 = QKDBuffer(key_size=256, max_capacity=10, min_capacity=2, ip_address="192.168.1.1")
    buffer1.monitor_buffer()
    buffer1.generate_and_add_key()
    buffer1.get_status()

    # Test adding multiple keys
    for i in range(5):
        key = QKDKey.generate_random_key(256)
        buffer1.add_key(key)

    # Test removing a key
    key_id_to_remove = list(buffer1.buffer.keys())[0]  # Get the first key ID
    buffer1.remove_key(key_id_to_remove)

    # Test getting a key
    key_id_to_get = list(buffer1.buffer.keys())[0]  # Get the first key ID
    key = buffer1.get_key(key_id_to_get)
    if key:
        print(f"Retrieved key with ID: {key.key_id}")

    # Test getting and removing a key
    key_id_to_get_remove = list(buffer1.buffer.keys())[0]  # Get the first key ID
    key = buffer1.get_remove_key(key_id_to_get_remove)
    if key:
        print(f"Retrieved and removed key with ID: {key.key_id}")

    # Test buffer status
    buffer1.monitor_buffer()

    # Print all keys in buffer
    buffer1.print_all_keys()
